
#在gpu上训练要将对应部分进行cuda（）方法转化
#数据、模型、损失函数可进行cuda（）方法转化,其余无法zhuan
#可以使用to（）方法指定设备
import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time
from model import *

# 准备数据集
train_data = torchvision.datasets.CIFAR10("datavision", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10("datavision", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
#载入数据集
train_dataloader = DataLoader(train_data,batch_size=64,shuffle=True)
test_dataloader = DataLoader(test_data,batch_size=64,shuffle=True)
#搭建神经网络
#常将神经网络模型在另一个文件中实现
MyModule = Module()


#模型的cuda转化
MyModule = MyModule.cuda()

'''
指定设备
device = torch.device("cuda")
MyModule = MyModule.to(device)
使用此方法可以类似直接指定cuda（）方法进行转化

'''
#设置损失函数、优化器
loss_func = nn.CrossEntropyLoss()
learning_rate = 1e-2
optimizer = torch.optim.SGD(MyModule.parameters(),lr=learning_rate)

#损失函数的转化
loss_func = loss_func.cuda()

#设置训练相关参数
train_step = 0
test_step = 0
step = 10

#设置tensorboard
writer = SummaryWriter("log6")

#训练模型
#数据在遍历中进行cuda（）转化
for epoch in range(step):
    print("-----------第{}轮训练-----------".format(epoch+1))
    start_time = time.time()
    #训练步骤
    MyModule.train()
    for imgs,targets in train_dataloader:
        imgs = imgs.cuda()
        targets = targets.cuda()
        outputs = MyModule(imgs)
        loss = loss_func(outputs,targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_step += 1
        if train_step % 100 == 0:
            end_time = time.time()
            print("\n---第{}次训练---\nloss: {}".format(train_step,loss.item()))
            print("训练时间：{}".format(end_time-start_time))
            writer.add_scalar("train_loss",loss,train_step)

    #测试步骤
    MyModule.eval()
    total_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for imgs,targets in test_dataloader:
            imgs = imgs.cuda()
            targets = targets.cuda()
            outputs = MyModule(imgs)
            loss = loss_func(outputs,targets)
            total_loss += loss.item()
            total_accuracy += (outputs.argmax(1) == targets).sum()
        test_step += 1
        print("测试集的正确率:{}".format(total_accuracy))
        writer.add_scalar("test_loss",total_loss,test_step)
    torch.save(MyModule,"MyModule_train{0}.pth".format(epoch))

writer.close()


